use std::ops::Sub;

use clap::Parser;
//use tokio::sync::mpsc;
use futures::channel::mpsc;
use futures::{sink::SinkExt, stream::StreamExt};
mod pprof;

#[derive(Parser, Clone, Debug, Default)]
#[clap(author, version, about, long_about = None)]
struct Opts {
    /// the size of event, no used when event type is int(0)
    #[clap(long, short, default_value = "16")]
    size: i64,

    /// event type, 0 for int 1 for str
    #[clap(long, short = 't', default_value = "0")]
    etype: i64,

    /// the count of workers
    #[clap(long, short, default_value = "100")]
    worker: usize,

    /// the count of events per worker
    #[clap(long, short, default_value = "100")]
    event: usize,

    /// the size of worker queue
    #[clap(long, short, default_value = "16")]
    queue: usize,

    /// output as csv format, default will be json
    #[clap(long, short)]
    csv: bool,

    /// cpuprofile result
    #[clap(long, default_value = "")]
    cpuprofile: String,

    /// more output
    #[clap(long, short)]
    verbose: bool,
}

//static SN : AtomicUsize = AtomicUsize::new(0);

#[derive(Debug, Clone, PartialEq, PartialOrd)]
enum Event {
    Int(i64),
    Str(String),
    CheapStr(&'static str),
}

impl Event {
    pub fn new(etype: i64, seed: i64) -> Self {
        match etype {
            0 => Self::Int(seed),
            1 => Self::Str("A".repeat(seed as usize)),
            2 => {
                let s = Box::new("A".repeat(seed as usize));
                let s: &'static str = Box::leak(s).as_str();
                Event::CheapStr(s)
            }
            3 => Self::Str("A".repeat(seed as usize)),
            _ => panic!("invalid event type"),
        }
    }
    pub fn is_exit(&self) -> bool {
        match self {
            Event::Int(v) => (-1).eq(v),
            Event::Str(v) => v.eq("exit"),
            Event::CheapStr(v) => "exit".eq(*v),
        }
    }
}

type EventSender = mpsc::Sender<Event>;
type EventReceiver = mpsc::Receiver<Event>;

async fn worker(mut queue: EventReceiver, mut done: mpsc::Sender<usize>, events: usize) {
    let mut n = 0;

    for _i in 0..events {
        if let Some(event) = queue.next().await {
            if !event.is_exit() {
                n += 1;
            }
        }
    }
    done.send(n).await.unwrap();
}

async fn dispatch_to(event: Event, events: usize, mut addr: EventSender) {
    for _ in 0..events {
        addr.feed(event.clone()).await.unwrap();
    }
    //addr.flush().await.unwrap();
    addr.close().await.unwrap();
}

async fn dispatch(opts: Opts, address: Vec<EventSender>) {
    let event = Event::new(opts.etype, opts.size);
    address.into_iter().for_each(|addr| {
        let event = event.clone();
        let events = opts.event;
        tokio::spawn(async move { dispatch_to(event, events, addr).await });
    });
}

#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let opts = Opts::parse();

    let prof = pprof::PProf::new(!opts.cpuprofile.is_empty());

    let (done_tx, mut done_rx) = mpsc::channel(opts.worker);
    let total_events = opts.event * opts.worker;

    // startup workers
    // address is the send port of all workers
    let mut address = vec![];
    for _ in 0..opts.worker {
        let (tx, rx) = mpsc::channel(opts.queue);
        address.push(tx);
        tokio::spawn(worker(rx, done_tx.clone(), opts.event));
    }

    // t1 is the time of startup
    let t1 = std::time::Instant::now();

    // spawn the dispatch
    tokio::spawn(dispatch(opts.clone(), address));

    // wait all task done
    let mut sn: usize = 0;
    for _i in 0..opts.worker {
        if let Some(n) = done_rx.next().await {
            sn += n;
        }
    }

    let t2 = std::time::Instant::now();
    if opts.verbose {
        // verfiy done is real done
        println!("total events {} done events {}", total_events, sn);
    }
    let ts = t2.sub(t1).as_secs_f64();
    let speed = (total_events) as f64 / ts;
    let etype = match opts.etype {
        0 => "int",
        1 => "str",
        2 => "str_ptr",
        3 => "str_clone",
        _ => "unknown",
    };
    if opts.csv {
        println!(
            "rust,{},{},{},{:.3},{:.0}",
            etype, opts.worker, opts.event, ts, speed
        )
    } else {
        println!("workers   : {}", opts.worker);
        println!("events    : {}", opts.event);
        println!("time used : {:.3}S", ts);
        println!("Speed     : {:.0}/S", speed);
    }
    if !opts.cpuprofile.is_empty() {
        prof.report(&opts.cpuprofile)?;
    }
    Ok(())
}
